from torch import nn


class MAELoss(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.cri = nn.L1Loss()

    def forward(self, pred, target):
        return self.cri(pred, target)


class MSELoss(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.cri = nn.MSELoss()

    def forward(self, pred, target):
        return self.cri(pred, target)


class MSE_MAE_Loss(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.cri_mse = nn.MSELoss()
        self.cri_mae = nn.L1Loss()

    def forward(self, pred, target):
        return self.cri_mse(pred, target) + self.cri_mae(pred, target)


class ChannelMaskedLoss(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.cri_mse = nn.MSELoss()

    def forward(self, pred_seq, target_seq,mask_pred,mask_target):
        l1 = self.cri_mse(pred_seq, target_seq)
        l2 = self.cri_mse(mask_pred, mask_target)
        return l1 + l2
